import os
import sys
import yaml


def get_real_labels(_, priority):
    real_labels = {}
    cmd = 'kubectl get nodes -o custom-columns=NODE:.metadata.name,AREA:.metadata.labels.area'
    results = os.popen(cmd)
    for r in results:
        if 'NODE' in r and 'AREA' in r:
            continue
        line_data = [i for i in r.split(' ') if i != '']
        node_uuid = line_data[0].strip()
        area_name = line_data[1].strip()
        if real_labels.get(area_name):
            real_labels[area_name].append(node_uuid)
        else:
            real_labels[area_name] = [node_uuid]

    for label in priority:
        if not real_labels.get(label):
            real_labels[label] = []

    return real_labels


def set_node_with_label(_, nodes, label):
    cmd = 'kubectl label nodes {} area={}'
    log = []
    for uuid in nodes:
        try:
            r = os.popen(cmd.format(uuid, label))
            for l in r:
                print('Label operator log: {}'.format(l))
            log.append('Succeed to set <{}> label for <{}> node.\n'.format(label, uuid))
        except Exception:
            log.append('Fail to set <{}> for <{}> node.\n'.format(label, uuid))

    return log


def get_expected_labels(file):
    with open(file) as f:
        clusters = yaml.safe_load(f)

    return clusters


def unify_none_nodes(real_labels):
    none_label = real_labels.pop('<none>') if real_labels.get('<none>') else []
    empty_label = real_labels.pop('') if real_labels.get('') else []
    if none_label or empty_label:
        real_labels[None] = none_label + empty_label
    else:
        real_labels[None] = []


def handle_mismatch_nodes(file, label, real_labels, expected_labels, priority, real_label_fun, set_node_with_label_fn):
    log = ['The current status is <{}> = <{}>, Expected status is <{}> = <{}>.\n'
           .format(label, len(real_labels.get(label)), label, expected_labels.get(label))]
    if real_labels.get(None):
        expected_node_number = expected_labels.get(label) - len(real_labels.get(label))
        log += set_node_with_label_fn('', real_labels.get(None)[:expected_node_number], label)
        real_labels[None] = real_labels.get(None)[expected_node_number:]
    else:
        log.append('No none node was found.\n')

    final_labels = real_label_fun(file[1], priority)
    log.append('The post-processing status is <{}> = <{}>, Expected status is <{}> = <{}>.\n'
               .format(label, len(final_labels.get(label)), label, expected_labels.get(label)))

    return log


def check_status(file, real_labels, expected_label, priority, real_label_fun, set_node_with_label_fn):
    log = []
    unify_none_nodes(real_labels)

    for label in priority:
        real_labels[label] = real_labels.get(label) if real_labels.get(label) else []
        if len(real_labels.get(label)) < expected_label.get(label):
            log += handle_mismatch_nodes(file, label, real_labels, expected_label, priority, real_label_fun,
                                         set_node_with_label_fn)
    if not log and real_labels.get(None):
        for node in real_labels.get(None):
            log.append('Find none node = <{}>.\n'.format(node))

    return log


def main(file, real_label_fun, set_node_with_label_fn):
    expected_data = get_expected_labels(file[0])
    real_labels = real_label_fun(file[1], expected_data['Priority'])

    log = check_status(file, real_labels, expected_data['Label'], expected_data['Priority'], real_label_fun,
                       set_node_with_label_fn)

    return log


if __name__ == '__main__':
    args = sys.argv
    if len(args) != 3:
        print('Usage: python check_nodes.py vault.yaml None.')
        exit(1)
    email_info = main(args[1:], get_real_labels, set_node_with_label)
    if email_info:
        print(email_info)
    else:
        print('successful')
